// mcts.cpp
#include "mcts.h"
#include <algorithm>
#include <random>
#include <memory>

// Implementation of MCTSNode methods

std::shared_ptr<MCTSNode> MCTSNode::add_random_child() {
    if (unvisited_moves.empty()) {
        return nullptr;
    }

    // Initialize random number generator
    static std::mt19937 static_rng(std::random_device{}());
    std::uniform_int_distribution<> dist(0, unvisited_moves.size() - 1);
    int index = dist(static_rng);

    // Select and remove the move from unvisited_moves
    Move new_move = unvisited_moves[index];
    unvisited_moves.erase(unvisited_moves.begin() + index);

    // Apply the move to get the new game state
    GameState new_game_state = game_state.apply_move(new_move);

    // Create a new child node
    std::shared_ptr<MCTSNode> child_node = std::make_shared<MCTSNode>(new_game_state, shared_from_this(), new_move);
    children.push_back(child_node);
    return child_node;
}

void MCTSNode::record_win(Player winner) {
    win_counts[winner] += 1;
    num_rollouts += 1;
}

bool MCTSNode::can_add_child() const {
    return !unvisited_moves.empty();
}

bool MCTSNode::is_terminal() const {
    return game_state.is_over();
}

double MCTSNode::winning_frac(Player player) const {
    if (num_rollouts == 0) return 0.0;
    return static_cast<double>(win_counts.at(player)) / static_cast<double>(num_rollouts);
}

// Implementation of MCTSAgent methods

Move MCTSAgent::select_move(const GameState& game_state) {
    // Initialize root node
    std::shared_ptr<MCTSNode> root = std::make_shared<MCTSNode>(game_state);

    // Run MCTS iterations
    for (int i = 0; i < num_rounds; ++i) {
        std::shared_ptr<MCTSNode> node = root;

        // Traverse the tree until reaching a node that can be expanded
        while (!node->can_add_child() && !node->is_terminal()) {
            node = select_child(node);
            if (!node) {
                break;
            }
        }

        // If node can be expanded, add a random child
        if (node && node->can_add_child()) {
            node = node->add_random_child();
        }

        // Simulate a random game from the node's game state
        if (node) {
            Player winner = simulate_random_game(node->game_state);
            // Backpropagate the result
            std::shared_ptr<MCTSNode> back_node = node;
            while (back_node) {
                back_node->record_win(winner);
                if (auto parent_ptr = back_node->parent.lock()) {
                    back_node = parent_ptr;
                }
                else {
                    back_node = nullptr;
                }
            }
        }

        // Optional: print progress
        std::cout << "winner " << i << std::endl;
    }

    // Choose the best child based on winning fraction
    double best_pct = -1.0;
    Move best_move;
    for (const auto& child : root->children) {
        double child_pct = child->winning_frac(game_state.next_player);
        if (child_pct > best_pct) {
            best_pct = child_pct;
            best_move = child->move;
        }
    }

    return best_move;
}

Player MCTSAgent::simulate_random_game(const GameState& game_state) {
    GameState current_state = game_state; // Assuming GameState has a copy constructor

    std::uniform_int_distribution<> move_dist;
    while (!current_state.is_over()) {
        std::vector<Move> possible_moves = current_state.legal_moves();
        if (possible_moves.empty()) {
            break;
        }
        // Select a random move
        std::uniform_int_distribution<> dist(0, possible_moves.size() - 1);
        int move_index = dist(rng);
        Move move = possible_moves[move_index];
        current_state = current_state.apply_move(move);
    }

    Player winner = current_state.get_winner();
    return winner;
}

std::shared_ptr<MCTSNode> MCTSAgent::select_child(std::shared_ptr<MCTSNode> node) {
    if (node->children.empty()) {
        return nullptr;
    }

    int total_rollouts = 0;
    for (const auto& child : node->children) {
        total_rollouts += child->num_rollouts;
    }

    if (total_rollouts == 0) {
        return nullptr;
    }

    double best_score = -1.0;
    std::shared_ptr<MCTSNode> best_child = nullptr;

    for (const auto& child : node->children) {
        double child_win_frac = child->winning_frac(node->game_state.next_player);
        double score = uct_score(total_rollouts, child->num_rollouts, child_win_frac, temperature);
        if (score > best_score) {
            best_score = score;
            best_child = child;
        }
    }

    return best_child;
}
